import os, time, math
import numpy as np
import pandas as pd
from multiprocessing import Pool
from kalman import execute_kalman_filter

def get_rotate_matrix_from(x: float, y: float) -> np.ndarray:
    theta = math.acos(-y/np.sqrt(x**2 + y**2))
    if x > 0:
        theta = 2*math.pi - theta
    cos_theta = math.cos(theta)
    sin_theta = math.sin(theta)
    return np.array([[cos_theta, -sin_theta], [sin_theta, cos_theta]])

def process_path(path: np.ndarray) -> np.ndarray:
    path = execute_kalman_filter(path)
    path: np.ndarray = path - path[19]
    rotate_matrix = get_rotate_matrix_from(*path[18])
    path = (rotate_matrix @ path.T).T
    path = path[19:][::5].reshape(-1,)
    return None if np.isnan(path).any() else path

def process_file(files, n):
    data, count = [], 0
    for file in files:
        df = pd.read_csv(os.path.join("data", file))
        df = df[df.OBJECT_TYPE == 'AGENT']
        path = df[['X', 'Y']].values
        data.append(process_path(path))
        count += 1
        if count % 1000 == 0: print(f"{n:3d}---Processing  {count} / {len(files)} files")
    data = [path for path in data if path is not None] 
    data = [path for path in data if path[3] > 0]
    data = np.array(data)
    np.savetxt(f'data-{n}.csv', data, delimiter=',')

def main():
    N_PROCESS = 16
    files = os.listdir("data")
    items = [[] for _ in range(N_PROCESS)]
    for file in files:
        name, _ = os.path.splitext(file)
        idx = int(name) % N_PROCESS
        items[idx].append(file)

    with Pool(N_PROCESS) as p:
        for i in range(N_PROCESS):
            p.apply_async(process_file, (items[i], i))
        p.close()
        p.join()

if __name__ == "__main__":
    s = time.time()
    main()
    e = time.time()
    print(f"Elapsed time: {e-s} seconds")